from argparse import ArgumentParser
import soundfile as sf
import torch
import torch.nn.functional as F
import numpy as np
import torch
import math
from tqdm import tqdm
from scipy.io import wavfile
from sklearn.metrics import roc_curve
from scipy.optimize import brentq
from scipy.interpolate import interp1d
from module.dataset import Evaluation_Dataset

def compute_dB(waveform):
    """
    Args:
        x (numpy.array): Input waveform (#length).
    Returns:
        numpy.array: Output array (#length).
    """
    val = max(0.0, np.mean(np.power(waveform, 2)))
    dB = 10*np.log10(val+1e-4)
    return dB

def load_audio(path, snr=25):
    sample_rate, waveform = wavfile.read(path)
    if snr>0:
        clean_dB = compute_dB(waveform)
        noise = np.random.randn(len(waveform))
        noise_dB = compute_dB(noise)
        noise = np.sqrt(10 ** ((clean_dB - noise_dB - snr) / 10)) * noise
        waveform = (waveform + noise)

    waveform = torch.from_numpy(waveform).unsqueeze(0).float()
    return waveform

def compute_eer(labels, scores):
    fpr, tpr, thresholds = roc_curve(labels, scores, pos_label=1)
    eer = brentq(lambda x: 1.0 - x - interp1d(fpr, tpr)(x), 0.0, 1.0)
    threshold = interp1d(fpr, thresholds)(eer)
    return eer, threshold

if __name__ == "__main__":
    from main import Task
    parser = ArgumentParser()
    parser = Task.add_model_specific_args(parser)
    parser.add_argument('--snr', type=int, default=-1)
    parser.add_argument('--top_n', type=int, help='trial file path', default=128)
    parser.add_argument('--checkpoint_path1', type=str, help='', default=None)
    parser.add_argument('--checkpoint_path2', type=str, help='', default=None)
    parser.add_argument('--cohort_path', type=str, help='', default=None)

    # model args
    args = parser.parse_args()

    state_dict1 = torch.load(args.checkpoint_path1, map_location="cpu")["state_dict"]
    state_dict2 = torch.load(args.checkpoint_path2, map_location="cpu")["state_dict"]

    for (k1, v1), (k2, v2) in zip(state_dict1.items(), state_dict2.items()):
        state_dict1[k1] = (v1+v2) * 0.5

    model = Task(**args.__dict__)
    model.load_state_dict(state_dict1, strict=True)
    model = model.eval()
    model = model.cuda()

    cohort = np.loadtxt(args.cohort_path, str)
    cohort_dataset = Evaluation_Dataset(cohort)
    loader = torch.utils.data.DataLoader(cohort_dataset, num_workers=10, shuffle=False, batch_size=1)

    model = model.eval()
    model = model.cuda()

    cohort_embeddings = []
    with torch.no_grad():
        for audio, audio_label in tqdm(loader):
            audio = audio.cuda()
            embedding = model(audio).detach().cpu().numpy()[0]
            cohort_embeddings.append(embedding)
    cohort_embeddings = np.array(cohort_embeddings)
    cohort_embeddings = (cohort_embeddings - np.mean(cohort_embeddings, axis=0)) / np.std(cohort_embeddings, axis=0)

    trials = np.loadtxt(args.trial_path, str)
    paths = np.unique(np.concatenate((trials.T[1], trials.T[2])))

    embs = []
    for idx, path in enumerate(tqdm(paths)):
        wavform = load_audio(path, args.snr)
        wavform = wavform.cuda()
        with torch.no_grad():
            embedding = model(wavform)
        embedding = embedding.detach().cpu().numpy()[0]
        embs.append(embedding)

    mean = np.mean(embs, axis=0)
    std = np.std(embs, axis=0)


    table = {}
    for idx, path in enumerate(tqdm(paths)):
        wavform = load_audio(path, args.snr)
        wavform = wavform.cuda()
        with torch.no_grad():
            embedding = model(wavform)
        embedding = embedding.detach().cpu().numpy()[0]
        embedding = (embedding - mean) / std

        scores = embedding.dot(cohort_embeddings.T)
        denoms = np.linalg.norm(embedding) * np.linalg.norm(cohort_embeddings, axis=1)
        scores = scores / denoms
        scores = np.sort(scores)
        val = []
        val.append(embedding)
        val.append(np.mean(scores[-args.top_n:]))
        val.append(np.std(scores[-args.top_n:]))
        table[path] = val

    scores = []
    labels = []
    for label, enroll_path, test_path in tqdm(trials):
        enroll_vector = table[enroll_path][0]
        enroll_mean = table[enroll_path][1]
        enroll_std = table[enroll_path][2]

        test_vector = table[test_path][0]
        test_mean = table[test_path][1]
        test_std = table[test_path][2]

        score = enroll_vector.dot(test_vector.T)
        denom = np.linalg.norm(enroll_vector) * np.linalg.norm(test_vector)
        score = score/denom

        score = 0.5 * (score-enroll_mean) / enroll_std + \
                0.5 * (score-test_mean) / test_std

        scores.append(score)
        labels.append(int(label))

    EER, threshold = compute_eer(labels, scores)
    print("cosine EER: {:.2f}% with threshold {:.2f}\n".format(EER*100, threshold))

